# !/usr/bin/env python
# -*- coding: utf-8 -*-

"""
速度限制中间件

令牌桶算法，不过这里做了简化,每个API请求需要的token数量都是统一的

"""

import os

from datetime import datetime
from starlette.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.db.redis_utils import red
from app.core.config import PROJECT_NAME

# 令牌桶大小
TOKEN_SIZE = int(os.getenv("API_LIMITER_SIZE", 90))
# 令牌生产速度
TOKEN_PRO = int(os.getenv("API_LIMITER_PRO", 10))
# 每个API需要的令牌数量
API_TOKEN_NUM = int(os.getenv("API_LIMITER_NUM", 30))


async def add2redis(ip: str, tk_num: int = TOKEN_SIZE):
    """
    把没有的IP加入red is中,并设置为当前时间，过期时间为60s
    """
    await red.execute('set', f'api:{PROJECT_NAME}_limiter_' + ip, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    await red.execute('set', f'api:{PROJECT_NAME}_limiter_current_token_num_' + ip, tk_num)
    await red.execute('EXPIRE', f'api:{PROJECT_NAME}_limiter_' + ip, 60)
    await red.execute('EXPIRE', f'api:{PROJECT_NAME}_limiter_current_token_num_' + ip, 60)


class LimiterMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        print('限流中间件')
        ip_or_host = request.client.host

        if 'X-Forwarded-For' in request.headers:
            ip_or_host = request.headers['X-Forwarded-For']

        if await red.execute("EXISTS", f"api:{PROJECT_NAME}_limiter_" + ip_or_host) == 0:
            print('IP没在red is限流中')
            # 不存在，就通过，并设置一个key存在redis中
            await add2redis(ip_or_host)
        else:
            # 存在，那就需要判断是否达到了限制条件
            pre_request_time = datetime.strptime(
                (await red.execute('get', f'api:{PROJECT_NAME}_limiter_' + ip_or_host)).decode(),
                "%Y-%m-%d %H:%M:%S")  # 上次请求时间
            print('上次访问时间', pre_request_time)
            now_request_time = datetime.now()  # 当前请求时间

            production_tokens_num = min((now_request_time - pre_request_time).seconds * TOKEN_PRO + \
                                        int((await red.execute('get',
                                                               f'api:{PROJECT_NAME}_limiter_current_token_num_' + ip_or_host)).decode()),
                                        TOKEN_SIZE)

            if production_tokens_num < API_TOKEN_NUM:
                # 达到了限制条件，那么token数量不减少
                await add2redis(ip_or_host, production_tokens_num)
                return JSONResponse("Your visit is too fast, please slow down.", status_code=401)
            # 如果没达到限制条件，那么token需要减去本次请求需要的token
            await add2redis(ip_or_host, production_tokens_num - API_TOKEN_NUM)

        response = await call_next(request)
        return response
